import os
import shutil
import subprocess
import yaml
from loguru import logger
from modelops_backend.settings import BASE_DIR


# 定义 GitUtils 类，用于管理 Git 仓库
class GitUtils:
    # 初始化方法，接受一个参数 model_name
    def __init__(self, model_name):
        # 初始化属性
        self.branch_list = []  # 存储分支列表
        self.model_name = model_name  # 存储模型名称
        self.base_path = "/data1/xyz/modelops-git"  # 存储 Git 仓库基础路径
        # 根据模型名称生成克隆路径
        # TODO:这里改成本地路径？
        # self.clone_path = "modelops-git@cc.qk0.cc:/home/modelops-git/{}/{}.git".format(self.model_name, self.model_name)
        self.clone_path = "xyz@162.105.88.63:22:/data1/xyz/modelops-git/{}/{}.git".format(self.model_name, self.model_name)
        self.current_branch = None  # 存储当前分支名称，默认为 None
        # 获取当前文件路径，并根据该路径生成 shell 脚本路径
        self.file_path = os.path.dirname(os.path.abspath(__file__))
        self.shell_path = os.path.join(self.file_path, "create_repo.sh")

    # 创建 Git 仓库方法
    def create_git_repo(self):
        # 判断文件夹是否存在
        pre_git_path = os.path.join(self.base_path, self.model_name)
        if os.path.exists(pre_git_path):
            # 如果文件夹已存在，则返回 False 和提示信息
            return False, ["Git文件夹已存在"]
        # 调用 subprocess_popen 方法执行 shell 脚本，创建 Git 仓库
        status, result = self.subprocess_popen("sh -x {} {}".format(self.shell_path, self.model_name))
        if not status:
            # 如果创建失败，则删除可能存在的文件夹
            if os.path.exists(pre_git_path):
                shutil.rmtree(pre_git_path)
        # 记录日志并返回执行结果
        logger.info("{} 创建Git仓库结果：{}".format(self.model_name, result))
        return status, result

    # 删除 Git 仓库方法
    def delete_git_repo(self):
        # 获取 Git 仓库路径
        pre_git_path = os.path.join(self.base_path, self.model_name)
        # 如果文件夹存在，则删除文件夹并返回删除成功的提示信息，否则返回文件夹不存在的提示信息
        if os.path.exists(pre_git_path):
            shutil.rmtree(pre_git_path)
            return True, ["Git文件夹删除成功"]
        else:
            return True, ["Git文件夹不存在"]

    # 判断 Git 仓库是否为空方法
    def is_empty(self):
        # 调用 get_all_branch 方法获取所有分支列表
        status, branch_list = self.get_all_branch()
        # 如果获取失败或者分支列表为空，则返回 True，表示仓库为空
        if status is False or branch_list == []:
            return True
        # 否则返回 False，表示仓库不为空
        return False

    # 获取 zip 文件路径方法
    def get_zip_path(self):
        # 获取 Git 仓库父路径和 zip 文件路径，并返回
        parent_path = os.path.join(self.base_path, self.model_name)
        file_path = os.path.join(self.base_path, self.model_name + ".zip")
        return parent_path, file_path

    # 获取仓库信息方法
    def get_repo_info(self):
        # 初始化返回数据字典
        return_data = dict()
        # 调用 get_all_branch 方法获取所有分支列表
        status, result = self.get_all_branch()
        # 如果获取失败，则直接返回失败结果
        if not status:
            return status, result
        # 否则设置当前分支和分支列表
        return_data['currrent_branch'] = "master" if self.current_branch is None else self.current_branch
        return_data['branch_list'] = ["master"] if len(self.branch_list) == 0 else self.branch_list
        # 获取文件列表、克隆路径、依赖信息和参数信息，并返回成功结果和数据字典
        return_data['file_list'] = self.list_files()
        return_data['clone_path'] = self.clone_path
        return_data['dependency'] = self.get_dependency_info()
        return_data['parameter'] = self.get_parameter_info()
        return True, return_data

    # 获取依赖信息方法
    def get_dependency_info(self):
        # 调用 get_file_content 方法获取 dependency.txt 文件内容
        status, file_content = self.get_file_content("dependency.txt")
        # 如果获取失败，则返回空字典
        if not status:
            return dict()
        # 否则解析文件内容为字典格式，并返回
        denpency = file_content.strip().split("\n")
        denency_dict = {key: value for key, value in [item.split("==") for item in denpency]}
        return denency_dict

    # 获取参数信息方法
    def get_parameter_info(self):
        # 获取配置文件路径
        pre_git_path = os.path.join(self.base_path, self.model_name)
        file_path = os.path.join(pre_git_path, "config.yaml")
        # 如果配置文件不存在，则返回空字典
        if not os.path.exists(file_path):
            return dict()
        # 否则读取配置文件内容，并返回
        with open(file_path, 'r', encoding="utf-8") as f:
            config = yaml.load(f, Loader=yaml.FullLoader)
        return config

    # 列出文件方法
    def list_files(self):
        # 获取 Git 仓库路径
        pre_git_path = os.path.join(self.base_path, self.model_name)
        file_list = []
        # 遍历文件夹下所有文件，排除 .git 文件夹，获取文件列表
        for file in os.listdir(pre_git_path):
            if '.git' in file:
                continue
            file_list.append(file)
        return file_list

    # 获取文件内容方法
    def get_file_content(self, file_path):
        try:
            # 获取文件路径
            pre_git_path = os.path.join(self.base_path, self.model_name)
            file_path = os.path.join(pre_git_path, file_path)
            # 如果文件不存在，则返回获取失败提示信息
            if not os.path.exists(file_path):
                return False, "文件不存在"
            # 否则读取文件内容，并返回获取成功和文件内容
            with open(file_path, 'r', encoding="utf-8") as f:
                content = f.read()
        except Exception as e:
            # 捕获异常，并返回获取失败提示信息
            logger.error("获取文件内容失败：{}".format(e))
            return False, e
        return True, content

    # 执行 git 命令方法
    def git_command(self, command):
        # 获取 Git 仓库路径
        pre_git_path = os.path.join(self.base_path, self.model_name)
        # 断言文件夹存在，否则抛出异常
        assert os.path.exists(pre_git_path), "Git文件夹{}不存在".format(pre_git_path)
        # 拼接真实的 git 命令，并执行
        real_command = "cd /data1/xyz/modelops-git/{};git --work-tree=. --git-dir={}.git {}".format(self.model_name,
                                                                                               self.model_name, command)
        # 调用 subprocess_popen 方法执行命令，并返回结果
        status, result = self.subprocess_popen(real_command)
        return status, result

    # 获取所有分支方法
    def get_all_branch(self):
        # 执行 git 命令获取所有分支列表
        status, result = self.git_command("branch -a")
        # 如果执行成功，则解析分支列表，设置当前分支和分支列表，最后返回结果
        if status:
            branch_list = []
            for branch in result:
                if '*' in branch:
                    self.current_branch = branch.split('*')[1].strip()
                    branch = branch.split('*')[1].strip()
                branch_list.append(branch)
            self.branch_list = branch_list
            return status, branch_list
        # 否则直接返回执行结果
        else:
            return status, result

    # 切换分支方法
    def checkout_branch(self, branch_name):
        # 执行 git 命令切换分支，并返回结果
        status, result = self.git_command("checkout {}".format(branch_name))
        return status, result

    # subprocess_popen 方法，用于执行 shell 命令
    def subprocess_popen(self, statement):
        # 执行 shell 命令，并定义输出格式
        p = subprocess.Popen(statement, shell=True, stdout=subprocess.PIPE)
        while p.poll() is None:  # 判断进程是否结束
            if p.wait() is not 0:  # 判断是否执行成功
                # 如果执行失败，则记录日志并返回失败结果
                logger.info("命令执行失败，请检查设备连接状态")
                return False, ["命令执行失败，请检查设备连接状态"]
            else:
                re = p.stdout.readlines()  # 获取原始执行结果
                result = []
                # 循环转换编码，并去除换行符
                for i in range(len(re)):
                    res = re[i].decode('utf-8').strip('\r\n').strip()
                    result.append(res)
                # 返回执行成功和结果
                return True, result

    # os_popen 方法，用于执行 shell 命令，并返回结果
    def os_popen(self, statement):
        re = os.popen(statement).readlines()
        result = []
        for i in range(0, len(re) - 1):  # 由于原始结果需要转换编码，所以循环转为utf8编码并且去除\n换行
            res = re[i].strip('\n')
            result.append(res)
        return result



if __name__ == '__main__':
    git = GitUtils("bert-softmax")
    # git.checkout_branch("dev")
    # logger.info(git.create_git_repo())
    # logger.info(git.list_files())
    # file_content = git.get_file_content("b.py")
    # print(file_content)
    print(git.get_all_branch())
    print(git.current_branch, git.branch_list)
    # print(git.checkout_branch("master"))
